{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e5e5f11a-519e-412d-b433-dfa5aeab5530",
   "metadata": {},
   "source": [
    "### SNP and indel enrichment "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4f599a6e-a079-4bf0-89a4-763a2160cdc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd \n",
    "from pybedtools import BedTool\n",
    "from pathlib import Path\n",
    "from io import StringIO\n",
    "import pysam\n",
    "import glob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "391def91-2fc2-4971-893f-cb08665f9c3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "wkdir = \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/misexpression_v3\"\n",
    "wkdir_path = Path(wkdir)\n",
    "# variables\n",
    "chrom=\"chr21\"\n",
    "ge_mtx_path = wkdir_path.joinpath(\"2_misexp_qc/misexp_gene_cov_corr/tpm_zscore_4568smpls_8610genes_flat_misexp_corr_qc.csv\")\n",
    "gt_info_root= wkdir_path.joinpath(\"4_vrnt_enrich/snp_indel_count_carriers/vrnts_gts_intersect\")\n",
    "genes_bed_path = wkdir_path.joinpath(f\"4_vrnt_enrich/snp_indel_count_carriers/genes_bed/{chrom}_genes.bed\")\n",
    "vrnts_bed_path = wkdir_path.joinpath(f\"4_vrnt_enrich/snp_indel_count_carriers/vrnts_bed/{chrom}_vrnts.bed\")\n",
    "gencode_path = wkdir_path.joinpath(\"reference/gencode.v31.annotation.sorted.gtf.gz\")\n",
    "vcf_path = f\"/lustre/scratch126/humgen/projects/interval_wgs/final_release_freeze/gt_phased/interval_wgs.{chrom}.gt_phased.vcf.gz\"\n",
    "wgs_rna_paired_smpls_path =wkdir_path.joinpath(\"1_rna_seq_qc/wgs_rna_match/paired_wgs_rna_postqc_prioritise_wgs.tsv\")\n",
    "out_dir =wkdir_path.joinpath(\"4_vrnt_enrich/snp_indel_count_carriers/tests/count_carriers\")\n",
    "\n",
    "z_cutoff_list = [2]\n",
    "window_start = 0\n",
    "gene_window_size = 10000\n",
    "window_step_size = 200000\n",
    "window_max = 1000000\n",
    "maf_range_list = [[0, 0.01], [0.01, 0.05], [0.05, 0.1], [0.1, 0.5]]\n",
    "tpm_cutoff = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b831a4f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of misexpressed genes on chromosome: 82\n",
      "Samples with genotyping data and passing RNA-seq QC: 2821\n",
      "upstream_0\n",
      "upstream_200000\n",
      "upstream_400000\n",
      "upstream_600000\n",
      "upstream_800000\n",
      "upstream_1000000\n",
      "downstream_0\n",
      "downstream_200000\n",
      "downstream_400000\n",
      "downstream_600000\n",
      "downstream_800000\n",
      "downstream_1000000\n"
     ]
    }
   ],
   "source": [
    "# create output directory\n",
    "out_dir_path = Path(out_dir)\n",
    "out_dir_path.mkdir(parents=True, exist_ok=True)\n",
    "# genotype information directory\n",
    "gt_info_root_path = Path(gt_info_root)\n",
    "# load gene expression matrix \n",
    "ge_mtx_df = pd.read_csv(ge_mtx_path)\n",
    "\n",
    "### subset gene expresstion matrix to genes on chromsome \n",
    "misexp_genes = pd.read_csv(genes_bed_path, sep=\"\\t\", header=None)[3].unique()\n",
    "num_misexp_genes = len(misexp_genes)\n",
    "print(f\"Number of misexpressed genes on chromosome: {num_misexp_genes}\")\n",
    "ge_mtx_chrom_df = ge_mtx_df[ge_mtx_df.gene_id.isin(misexp_genes)]\n",
    "\n",
    "### subset to gene expression matrix to samples with genotyping data \n",
    "# get EGAN IDs with RNA data \n",
    "wgs_rna_paired_smpls_df = pd.read_csv(wgs_rna_paired_smpls_path, sep=\"\\t\")\n",
    "egan_ids_with_rna = wgs_rna_paired_smpls_df.egan_id.unique()\n",
    "# load vcf\n",
    "vcf = pysam.VariantFile(vcf_path, mode = \"r\")\n",
    "vcf_samples = [sample for sample in vcf.header.samples]\n",
    "vcf_egan_ids_with_rna = set(egan_ids_with_rna).intersection(set(vcf_samples))\n",
    "# subset egan ID and RNA ID links to samples with genotyping data and passing QC \n",
    "wgs_rna_paired_smpls_in_vcf_df = wgs_rna_paired_smpls_df[wgs_rna_paired_smpls_df.egan_id.isin(vcf_egan_ids_with_rna)]\n",
    "ge_matrix_flat_chrom_egan_df = pd.merge(ge_mtx_chrom_df, wgs_rna_paired_smpls_in_vcf_df, how=\"inner\", on=\"rna_id\")\n",
    "# add gene sample pairs \n",
    "ge_matrix_flat_chrom_egan_df[\"gene_smpl_pair\"] = ge_matrix_flat_chrom_egan_df.gene_id + \",\" + ge_matrix_flat_chrom_egan_df.egan_id\n",
    "\n",
    "rna_ids_pass_gt_rna_qc = ge_matrix_flat_chrom_egan_df.rna_id.unique()\n",
    "gene_ids_pass_qc = ge_matrix_flat_chrom_egan_df.gene_id.unique()\n",
    "print(f\"Samples with genotyping data and passing RNA-seq QC: {len(rna_ids_pass_gt_rna_qc)}\")\n",
    "if ge_matrix_flat_chrom_egan_df.shape[0] != num_misexp_genes * len(rna_ids_pass_gt_rna_qc): \n",
    "    raise ValueError(\"Number of genes and samples passsing QC does not match gene expression matrix size.\")\n",
    "\n",
    "### get variants in windows \n",
    "# load variant and gene bed files \n",
    "vrnts_bed = BedTool(vrnts_bed_path)\n",
    "genes_bed = BedTool(genes_bed_path)\n",
    "# column names for variant window intersection bed file \n",
    "intersect_bed_columns={0:\"chrom_gene\", 1:\"start_gene\", 2:\"end_gene\", 3: \"gene_id\", \n",
    "                       4:\"0\", 5: \"strand\", 6:\"chrom_vrnt\", 7:\"start_vrnt\", 8:\"end_vrnt\", 9:\"vrnt_id\", 10: \"AF\"}\n",
    "\n",
    "vrnts_in_windows_dict = {}\n",
    "for direction in [\"upstream\", \"downstream\"]: \n",
    "    window_size = window_start\n",
    "    seen_vrnt_gene_pairs = set()\n",
    "    while window_size <= window_max:\n",
    "        print(f\"{direction}_{window_size}\")\n",
    "        if direction == \"upstream\": \n",
    "            l, r = window_size, 0 \n",
    "        else: \n",
    "            l, r = 0, window_size\n",
    "        intersect_bed_str = StringIO(str(genes_bed.window(vrnts_bed, l=l, r=r,sw=True)))\n",
    "        intersect_bed_df = pd.read_csv(intersect_bed_str, sep=\"\\t\", header=None).rename(columns=intersect_bed_columns)\n",
    "        # variant gene ID column \n",
    "        intersect_bed_df[\"vrnt_gene_id\"] = intersect_bed_df[\"vrnt_id\"].astype(str) + \",\" + intersect_bed_df[\"gene_id\"].astype(str)\n",
    "        # subset to variant gene pairs that have not been seen in previous windows    \n",
    "        vrnt_gene_pairs_in_window = set(intersect_bed_df.vrnt_gene_id.unique())\n",
    "        new_vrnt_gene_pairs = vrnt_gene_pairs_in_window - seen_vrnt_gene_pairs\n",
    "        intersect_bed_new_vrnt_gene_df = intersect_bed_df[intersect_bed_df[\"vrnt_gene_id\"].isin(new_vrnt_gene_pairs)]\n",
    "        # add variant gene pairs to seen set \n",
    "        seen_vrnt_gene_pairs = seen_vrnt_gene_pairs.union(vrnt_gene_pairs_in_window)\n",
    "        vrnts_in_windows_dict[f\"{direction}_{window_size}\"] = intersect_bed_new_vrnt_gene_df\n",
    "        window_size += window_step_size\n",
    "\n",
    "# check upstream zero and downstream 0 are the same \n",
    "if not vrnts_in_windows_dict[\"upstream_0\"].equals(vrnts_in_windows_dict[\"downstream_0\"]): \n",
    "    print(f\"{chrom}: upstream and downstream gene body windows are not identical.\")\n",
    "\n",
    "intersect_gene_window_str = StringIO(str(genes_bed.window(vrnts_bed, w=gene_window_size)))\n",
    "intersect_gene_window_bed_df = pd.read_csv(intersect_gene_window_str, sep=\"\\t\", header=None).rename(columns=intersect_bed_columns)\n",
    "vrnts_in_windows_dict[f\"gene_body_window_{gene_window_size}\"] = intersect_gene_window_bed_df\n",
    "\n",
    "# create single window for gene body\n",
    "vrnts_in_windows_dict[\"gene_body\"] = vrnts_in_windows_dict.pop(\"upstream_0\")\n",
    "del vrnts_in_windows_dict[\"downstream_0\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a3b34ccc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ENSG00000206102\n"
     ]
    }
   ],
   "source": [
    "count = 0\n",
    "count_carriers = {}\n",
    "for z_cutoff in z_cutoff_list: \n",
    "    # select misexpression events \n",
    "    misexp_df = ge_matrix_flat_chrom_egan_df[(ge_matrix_flat_chrom_egan_df[\"z-score\"] > z_cutoff) & \n",
    "                                             (ge_matrix_flat_chrom_egan_df.TPM > tpm_cutoff)\n",
    "                                                     ]\n",
    "    # get gene-sample pairs passing cutoff \n",
    "    test_gene_ids = misexp_df.gene_id.unique()\n",
    "    test_gene_smpl_pairs = misexp_df.gene_smpl_pair.unique()\n",
    "    # get gene-sample pairs in control group, limit to genes with a misexpression event passing cutoff\n",
    "    cntrl_gene_smpl_pairs = ge_matrix_flat_chrom_egan_df[~ge_matrix_flat_chrom_egan_df.gene_smpl_pair.isin(test_gene_smpl_pairs) &\n",
    "                                                         ge_matrix_flat_chrom_egan_df.gene_id.isin(test_gene_ids)\n",
    "                                                        ].gene_smpl_pair.unique()\n",
    "    # check overlap between sets is empty \n",
    "    if len(set(test_gene_smpl_pairs).intersection(set(cntrl_gene_smpl_pairs))) != 0:\n",
    "        raise ValueError(\"Overlap between control and test gene-sample pair sets\")\n",
    "    # check number of test and control gene-pairs matches expected \n",
    "    total_test_control_pairs = len(test_gene_smpl_pairs) + len(cntrl_gene_smpl_pairs)\n",
    "    test_genes_by_smpls = len(test_gene_ids) * len(rna_ids_pass_gt_rna_qc)\n",
    "    if total_test_control_pairs !=  test_genes_by_smpls: \n",
    "        raise ValueError(\"Number of test and control gene pairs does not match test gene IDs by samples\")\n",
    "    for gene_id in test_gene_ids: \n",
    "        print(gene_id)\n",
    "        test_smpls = set([gene_smpl_pair.split(\",\")[1] for gene_smpl_pair in test_gene_smpl_pairs if f\"{gene_id},\" in gene_smpl_pair])\n",
    "        cntrl_smpls = set([gene_smpl_pair.split(\",\")[1] for gene_smpl_pair in cntrl_gene_smpl_pairs if f\"{gene_id},\" in gene_smpl_pair])\n",
    "        # get genotype info\n",
    "        gene_id_has_gt_intersect_path = [gene_id_path for gene_id_path in gt_info_root_path.glob(f\"{chrom}/*\") if gene_id in str(gene_id_path)]\n",
    "        # check if gene has variants in window\n",
    "        if len(gene_id_has_gt_intersect_path) == 0:\n",
    "            for window in vrnts_in_windows_dict.keys():\n",
    "                for vrnt_type in [\"snp\", \"indel\"]:\n",
    "                    for maf_range in maf_range_list:  \n",
    "                        count_carriers[count] = [z_cutoff, gene_id, window, vrnt_type, f\"{maf_range[0]}-{maf_range[1]}\", 0, len(test_smpls), 0, len(cntrl_smpls)]\n",
    "                        count += 1 \n",
    "        elif len(gene_id_has_gt_intersect_path) == 1: \n",
    "            gt_info_path = gt_info_root_path.joinpath(f\"{chrom}/{chrom}_{gene_id}_vrnts_gts_intersect.tsv\")\n",
    "            gt_info_df = pd.read_csv(gt_info_path, sep=\"\\t\")\n",
    "            for window in vrnts_in_windows_dict.keys():\n",
    "                window_df = vrnts_in_windows_dict[window]\n",
    "                vrnt_id_in_gene_window = set(window_df[window_df.gene_id == gene_id].vrnt_id.unique())\n",
    "                vrnt_id_in_gene_window_df = gt_info_df[gt_info_df.vrnt_id.isin(vrnt_id_in_gene_window)]\n",
    "                for vrnt_type in [\"snp\", \"indel\"]:\n",
    "                    vrnt_type_in_gene_window_df = vrnt_id_in_gene_window_df[vrnt_id_in_gene_window_df.vrnt_type == vrnt_type]\n",
    "                    for maf_range in maf_range_list: \n",
    "                        maf_lower, maf_upper = maf_range\n",
    "                        vrnt_ids_low_af = vrnt_type_in_gene_window_df[(vrnt_type_in_gene_window_df.AF >= maf_lower) &\n",
    "                                                                    (vrnt_type_in_gene_window_df.AF < maf_upper) \n",
    "                                                                   ].vrnt_id.unique()\n",
    "                        vrnt_ids_high_af = vrnt_type_in_gene_window_df[(vrnt_type_in_gene_window_df.AF <= (1 - maf_lower)) &\n",
    "                                                                     (vrnt_type_in_gene_window_df.AF > (1 - maf_upper)) \n",
    "                                                                    ].vrnt_id.unique()\n",
    "                        vrnt_ids_in_gene_window = set(vrnt_ids_low_af).union(set(vrnt_ids_high_af))\n",
    "                        # get list of carriers \n",
    "                        vrnt_id_carriers_df = vrnt_id_in_gene_window_df[vrnt_id_in_gene_window_df.vrnt_id.isin(vrnt_ids_in_gene_window)]\n",
    "                        vrnts_id_carriers = set(vrnt_id_carriers_df.vrnt_id.unique()) \n",
    "                        if vrnts_id_carriers != vrnt_ids_in_gene_window: \n",
    "                            raise ValueError(f\"Missing variant IDs from genotype file: {vrnts_id_carriers - vrnt_ids_in_gene_window}\")\n",
    "                        carriers = vrnt_id_carriers_df.egan_id.unique()\n",
    "                        # get test and control carriers \n",
    "                        test_carriers = test_smpls.intersection(carriers)\n",
    "                        cntrl_carriers = cntrl_smpls.intersection(carriers)\n",
    "                        count_carriers[count] = [z_cutoff, gene_id, window, vrnt_type, f\"{maf_range[0]}-{maf_range[1]}\", len(test_carriers), len(test_smpls), len(cntrl_carriers), len(cntrl_smpls)]\n",
    "                        count += 1 \n",
    "        else: \n",
    "            raise ValueError(f\"{gene_id} has multiple file names: {gene_id_has_gt_intersect_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "433e3701",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of misexpressed genes on chromosome: 82\n",
      "Samples with genotyping data and passing RNA-seq QC: 2821\n",
      "upstream_0\n",
      "upstream_200000\n",
      "upstream_400000\n",
      "upstream_600000\n",
      "upstream_800000\n",
      "upstream_1000000\n",
      "downstream_0\n",
      "downstream_200000\n",
      "downstream_400000\n",
      "downstream_600000\n",
      "downstream_800000\n",
      "downstream_1000000\n",
      "ENSG00000206102\n",
      "ENSG00000224100\n",
      "ENSG00000229289\n",
      "ENSG00000229986\n",
      "ENSG00000274248\n",
      "ENSG00000237338\n",
      "ENSG00000229382\n",
      "ENSG00000142182\n",
      "ENSG00000232360\n",
      "ENSG00000182591\n",
      "ENSG00000227702\n",
      "ENSG00000236471\n",
      "ENSG00000224541\n",
      "ENSG00000233393\n",
      "ENSG00000275874\n",
      "ENSG00000184856\n",
      "ENSG00000226956\n",
      "ENSG00000225330\n",
      "ENSG00000237735\n",
      "ENSG00000205439\n",
      "ENSG00000241123\n",
      "ENSG00000280095\n",
      "ENSG00000223400\n",
      "ENSG00000274749\n",
      "ENSG00000272804\n",
      "ENSG00000188694\n",
      "ENSG00000224574\n",
      "ENSG00000230379\n",
      "ENSG00000236545\n",
      "ENSG00000187766\n",
      "ENSG00000231324\n",
      "ENSG00000187026\n",
      "ENSG00000267857\n",
      "ENSG00000231620\n",
      "ENSG00000227075\n",
      "ENSG00000231986\n",
      "ENSG00000231106\n",
      "ENSG00000224269\n",
      "ENSG00000206105\n",
      "ENSG00000223608\n",
      "ENSG00000237664\n",
      "ENSG00000259981\n",
      "ENSG00000273115\n",
      "ENSG00000198390\n",
      "ENSG00000230978\n",
      "ENSG00000261706\n",
      "ENSG00000233480\n",
      "ENSG00000277693\n",
      "ENSG00000198054\n",
      "ENSG00000236332\n",
      "ENSG00000183640\n",
      "ENSG00000224413\n",
      "ENSG00000230794\n",
      "ENSG00000231867\n",
      "ENSG00000160202\n",
      "ENSG00000235123\n",
      "ENSG00000279579\n",
      "ENSG00000223563\n",
      "ENSG00000160181\n",
      "ENSG00000160182\n",
      "ENSG00000187175\n",
      "ENSG00000227342\n",
      "ENSG00000232806\n",
      "ENSG00000186977\n",
      "ENSG00000233316\n",
      "ENSG00000184032\n",
      "ENSG00000230972\n",
      "ENSG00000221864\n",
      "ENSG00000225637\n",
      "ENSG00000281181\n",
      "ENSG00000225298\n",
      "ENSG00000226935\n"
     ]
    }
   ],
   "source": [
    "# write to output\n",
    "output_columns=[\"z_cutoff\", \"gene_id\", \"window\", \"vrnt_type\", \"maf_range\", \"misexp_carrier\", \"misexp_total\", \"control_carrier\", \"control_total\"]\n",
    "count_carriers_df = pd.DataFrame.from_dict(count_carriers, orient=\"index\", columns=output_columns)\n",
    "path_out = out_dir_path.joinpath(f\"{chrom}_carrier_count.tsv\")\n",
    "count_carriers_df.to_csv(path_out, sep=\"\\t\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a8d0b0b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
